import re
import requests
import json
from urllib.parse import urlparse, parse_qs, unquote
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import time
import os

def extract_url(text):
    """
    从文本中提取第一个URL链接
    
    Args:
        text (str): 包含URL的文本
        
    Returns:
        str: 提取到的第一个URL，如果没有找到则返回None
    """
    # 使用正则表达式匹配URL
    url_pattern = r'https?://[^\s]+'
    urls = re.findall(url_pattern, text)
    return urls[0] if urls else None

def download_video(video_url, output_dir='downloads'):
    """
    下载视频文件
    
    Args:
        video_url (str): 视频下载链接
        output_dir (str): 输出目录
        
    Returns:
        str: 下载的视频文件路径，如果下载失败则返回None
    """
    try:
        # 创建输出目录
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            
        # 设置请求头
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36',
            'Referer': 'https://www.douyin.com/'
        }
        
        # 发送请求获取视频
        response = requests.get(video_url, headers=headers, stream=True)
        response.raise_for_status()
        
        # 生成文件名
        filename = f"douyin_video_{int(time.time())}.mp4"
        filepath = os.path.join(output_dir, filename)
        
        # 下载视频
        print(f"正在下载视频到: {filepath}")
        with open(filepath, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
                    
        print("视频下载完成！")
        return filepath
        
    except Exception as e:
        print(f"下载视频失败: {str(e)}")
        return None

def get_video_url(share_url):
    """
    获取抖音分享链接中的视频真实下载链接
    
    Args:
        share_url (str): 抖音分享链接
        
    Returns:
        str: 视频真实下载链接，如果获取失败则返回None
    """
    try:
        # 设置Chrome选项
        chrome_options = Options()
        chrome_options.add_argument('--headless')  # 无头模式
        chrome_options.add_argument('--disable-gpu')
        chrome_options.add_argument('--no-sandbox')
        chrome_options.add_argument('--disable-dev-shm-usage')
        chrome_options.add_argument('--user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/135.0.0.0 Safari/537.36')
        
        # 设置ChromeDriver路径
        chromedriver_path = os.path.join(os.getcwd(), 'chromedriver', 'chromedriver-win64', 'chromedriver.exe')
        print(f"ChromeDriver路径: {chromedriver_path}")
        
        if not os.path.exists(chromedriver_path):
            raise FileNotFoundError(f"ChromeDriver not found at: {chromedriver_path}")
            
        service = Service(executable_path=chromedriver_path)
        
        # 初始化浏览器
        driver = webdriver.Chrome(service=service, options=chrome_options)
        
        try:
            # 访问页面
            print("正在访问页面...")
            driver.get(share_url)
            
            # 等待页面加载
            print("等待页面加载...")
            time.sleep(5)  # 等待页面完全加载
            
            # 获取页面源码
            print("获取页面源码...")
            page_source = driver.page_source
            
            # 从页面内容中提取视频链接
            patterns = [
                r'https://v3-web\.douyinvod\.com/[^\s"\']+',  # 直接匹配视频链接
                r'play_addr.*?url_list.*?\["(https://[^"]+)"',  # 匹配JSON中的视频链接
                r'video_url.*?"(https://[^"]+)"',  # 匹配视频URL
                r'play_addr.*?url.*?"(https://[^"]+)"'  # 匹配播放地址
            ]
            
            print("正在提取视频链接...")
            for pattern in patterns:
                matches = re.findall(pattern, page_source)
                if matches:
                    # 获取第一个匹配的链接
                    video_url = matches[0]
                    # 清理链接（移除可能的转义字符）
                    video_url = video_url.replace('\\u0026', '&')
                    return video_url
            
            return None
            
        finally:
            # 关闭浏览器
            print("关闭浏览器...")
            driver.quit()
            
    except Exception as e:
        print(f"获取视频链接失败: {str(e)}")
        return None

def main():
    """
    主函数，处理用户输入并执行下载
    """
    print("欢迎使用抖音视频下载工具！")
    print("请输入包含抖音分享链接的文本（可以直接粘贴）：")
    print("（输入 'q' 或 'quit' 退出程序）")
    
    while True:
        text = input("\n请输入文本: ").strip()
        
        # 检查是否要退出
        if text.lower() in ['q', 'quit']:
            print("程序已退出。")
            break
            
        if not text:
            print("输入不能为空，请重新输入！")
            continue
            
        # 提取并处理链接
        url = extract_url(text)
        if not url:
            print("未找到有效的抖音分享链接，请重新输入！")
            continue
            
        print("\n提取到的分享链接:")
        print(url)
        
        video_url = get_video_url(url)
        if not video_url:
            print("获取视频链接失败，请重试！")
            continue
            
        print("\n视频真实下载链接:")
        print(video_url)
        
        # 下载视频
        download_video(video_url)

# 运行主函数
if __name__ == "__main__":
    main()